#include <Uefi.h>
#include <Base.h>
#include <IndustryStandard/Pci.h>
#include <Library/BaseLib.h>
#include <Library/UefiLib.h>
#include <Library/PrintLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiApplicationEntryPoint.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/IoLib.h>
#include "DRAMApp.h"
#include "JEP106.h"

UINT16      gSmbusIOBase = 0;
DDR4_INFO   gDRAMInfo[MAX_DRAM_SLOT_NUM];
DDR5_INFO   gDDR5DRAMInfo[MAX_DRAM_SLOT_NUM];
UINT8       DRAMSlaveAddress[MAX_DRAM_SLOT_NUM] = {
              0xA0, // 0000
              0xA2, // 0010
              0xA4, // 0100
              0xA6, // 0110
              0xA8, // 1000
              0xAA, // 1010
              0xAC, // 1100
              0xAE  // 1110
            };

UINT8
IOCF8Read8 (
  UINT32  Address
  )
{
  IoWrite32 (PCI_IO_CF8, Address);
  return IoRead8 (PCI_IO_CFC + (Address & 0x03));
}

UINT16
IOCF8Read16 (
  UINT32  Address
  )
{
  IoWrite32 (PCI_IO_CF8, Address);
  return IoRead16 (PCI_IO_CFC + (Address & 0x02));
}

UINT32
IOCF8Read32 (
  UINT32  Address
  )
{
  IoWrite32 (PCI_IO_CF8, Address);
  return IoRead32 (PCI_IO_CFC);
}

EFI_STATUS
GetSmbusIoPortByIO(
  OUT UINT16  *IoPort
  )
{
  EFI_STATUS  Status = EFI_NOT_FOUND;
  UINT16      BusIndex   = 0;
  UINT16      DevIndex   = 0;
  UINT16      FunIndex   = 0;
  UINT16      VendorID   = 0xFFFF;
  UINT32      Data32     = 0;
  UINT8       HeaderType = 0;
  UINT32      ClassBuffer = 0;
  UINT16      Command     = 0;
  UINT8       Index   = 0;
  UINT8       Offset  = 0;
  UINT32      BarData = 0;
  UINT8       FindFlag = 0;
  UINT8       Bar5Flag = 0;

  for (BusIndex = 0; BusIndex < MAX_BUS_NUM; BusIndex++) {
    for (DevIndex = 0; DevIndex < MAX_DEV_NUM; DevIndex++) {
      for (FunIndex = 0; FunIndex < MAX_FUN_NUM; FunIndex++) {
        VendorID = IOCF8Read16(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_VENDOR_ID_OFFSET);
        if (VendorID == 0xFFFF) {
          if (FunIndex == 0) {
            break;
          }
          continue;
        }

        ClassBuffer = IOCF8Read32(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_REVISION_ID_OFFSET);
        if ((((ClassBuffer >> 8) & 0xFF) == 0x00) &&
            (((ClassBuffer >> 16) & 0xFF) == 0x05) &&
            (((ClassBuffer >> 24) & 0xFF) == 0x0C)) { // smbus controller
          Command = IOCF8Read16(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_COMMAND_OFFSET);
          if ((Command & BIT0) == BIT0) { // IO Enable
            for (Index = 0; Index < 6; Index++) {
              Offset = (Index * 4);
              BarData = IOCF8Read32(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_BASE_ADDRESSREG_OFFSET + Offset);
              if ((BarData & 0x01) == 1) { // IO
                if ((BarData & 0xFFF0) != 0) {
                  FindFlag++;
                  if (Offset == 0x10) {
                    Bar5Flag = 1;
                  }
                  *IoPort = (UINT16)(BarData & 0xFFF0);
                }
              }
            }

            if ((FindFlag > 1) && (Bar5Flag == 1)) { // Fing more bar
              BarData = IOCF8Read32(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_BASE_ADDRESSREG_OFFSET + 0x10);
              *IoPort = (UINT16)(BarData & 0xFFF0);
            }

            if (FindFlag == 0) { // NOt IO Mapping BAR
              return EFI_ABORTED;
            }

            return EFI_SUCCESS;
          } else { // IO Not Enable
            return EFI_ABORTED;
          }
        }

        Data32 = IOCF8Read32(PCI_IO_CF8_ADDRESS(BusIndex, DevIndex, FunIndex) + PCI_CACHELINE_SIZE_OFFSET);
        HeaderType = (Data32 >> 16) & 0xFF;
        if ((FunIndex == 0) && ((HeaderType & HEADER_TYPE_MULTI_FUNCTION) == 0)) {
          break;
        }
      }
    }
  }

  return Status;
}

EFI_STATUS
GetSmbusIOPortByMmio (
  OUT UINT16 *IoPort
  )
{
  EFI_STATUS  Status = EFI_SUCCESS;
  UINT8       BusNum = 0;
  UINT8       DevNum = 0;
  UINT8       FunNum = 0;
  UINT16      VendorID   = 0xFFFF;
  UINT8       HeaderType = 0;
  UINT32      ClassBuffer = 0;
  UINT8       Index   = 0;
  UINT8       Offset  = 0;
  UINT32      BarData = 0;

  for (BusNum = 0; BusNum < MAX_BUS_NUM; BusNum++) {
    for (DevNum = 0; DevNum < MAX_DEV_NUM; DevNum++) {
      for (FunNum = 0; FunNum < MAX_FUN_NUM; FunNum++) {
        VendorID = MmioRead16 (PCI_REG_BASE(BusNum, DevNum, FunNum, PCI_VENDOR_ID_OFFSET));
        if (VendorID == 0xFFFF) {
          if (FunNum == 0) {
            break;
          }
          continue;
        }

        ClassBuffer = MmioRead32 (PCI_REG_BASE(BusNum, DevNum, FunNum, PCI_REVISION_ID_OFFSET));
        if ((((ClassBuffer >> 8) & 0xFF) == 0x00) &&
            (((ClassBuffer >> 16) & 0xFF) == 0x05) &&
            (((ClassBuffer >> 24) & 0xFF) == 0x0C) ) {
          for (Index = 0; Index < 6; Index++) {
            Offset = (Index * 4);
            BarData = MmioRead32 (PCI_REG_BASE(BusNum, DevNum, FunNum, PCI_BASE_ADDRESSREG_OFFSET + Offset));
            if ((BarData & 0x01) == 1) { // IO
              if ((BarData & 0xFFF0) == 0) {
                return EFI_ABORTED;
              }
              *IoPort = (UINT16)(BarData & 0xFFF0);
              return EFI_SUCCESS;
            }
          }
        }

        HeaderType = MmioRead8 (PCI_REG_BASE(BusNum, DevNum, FunNum, PCI_HEADER_TYPE_OFFSET));
        if ((FunNum == 0) && ((HeaderType & HEADER_TYPE_MULTI_FUNCTION) == 0)) {
          break;
        }
      }
    }
  }

  return Status;
}

EFI_STATUS
ReadByteBySmbus (
  IN   UINT8  SlaveAddress,
  IN   UINT8  Command,
  IN   UINT8  Offset,
  OUT  UINT8  *Vaule
  )
{
  EFI_STATUS  Status = EFI_SUCCESS;
  UINT8   HostStatus = 0;
  UINT8   Retry = RETRY_COUNT;
  UINT8   Address = 0;

  if (gSmbusIOBase == 0xFF) {
    return EFI_INVALID_PARAMETER;
  }

  while (Retry != 0) {
    HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    while ((HostStatus & BIT0) == BIT0) { // smbus host busy
      HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    }

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR); // clear

    Address = (UINT8)((SlaveAddress << 1) | BIT0); // BIT0: Read
    IoWrite8 (gSmbusIOBase + SMBUS_HOST_ADR, Address);

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_CMD, Offset);

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_CTL, (UINT8)(Command | SMBUS_HOST_CTL_CMD_START));

    HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    while ((HostStatus & BIT0) == BIT0) {
      HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    }

    if (((HostStatus & BIT7) == BIT7) || 
        ((HostStatus & BIT4) == BIT4) ||
        ((HostStatus & BIT3) == BIT3) ||
        ((HostStatus & BIT2) == BIT2)) {
      Status = EFI_DEVICE_ERROR;
      IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR);
      Retry--;
      continue;
    } else {
      *Vaule = IoRead8 (gSmbusIOBase + SMBUS_HOST_DATA_0);
      Status = EFI_SUCCESS;
      break;
    }
  }

  IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR);

  return Status;
}


EFI_STATUS
WriteByteBySmbus (
  IN   UINT8  SlaveAddress,
  IN   UINT8  Command,
  IN   UINT8  Offset,
  IN   UINT8  Vaule
  )
{
  EFI_STATUS  Status = EFI_SUCCESS;
  UINT8   HostStatus = 0;
  UINT8   Retry = RETRY_COUNT;
  UINT8   Address = 0;

  if (gSmbusIOBase == 0xFF) {
    return EFI_INVALID_PARAMETER;
  }

  while (Retry != 0) {
    HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    while ((HostStatus & BIT0) == BIT0) { // smbus host busy
      HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    }

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR); // clear

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_DATA_0, Vaule);

    Address = (UINT8)((SlaveAddress << 1) | 0x00); // Write
    IoWrite8 (gSmbusIOBase + SMBUS_HOST_ADR, Address);

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_CMD, Offset);

    IoWrite8 (gSmbusIOBase + SMBUS_HOST_CTL, (UINT8)(Command | SMBUS_HOST_CTL_CMD_START));

    HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    while ((HostStatus & BIT0) == BIT0) {
      HostStatus = IoRead8 (gSmbusIOBase + SMBUS_HOST_STATUS);
    }

    if (((HostStatus & BIT7) == BIT7) ||
        ((HostStatus & BIT4) == BIT4) ||
        ((HostStatus & BIT3) == BIT3) ||
        ((HostStatus & BIT2) == BIT2)) {
      Status = EFI_DEVICE_ERROR;
      IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR);
      Retry--;
      continue;
    } else {
      Status = EFI_SUCCESS;
      break;
    }
  }

  IoWrite8 (gSmbusIOBase + SMBUS_HOST_STATUS, SMBUS_HOST_STA_CLEAR);

  return Status;
}

BOOLEAN
IsSPD5 ()
{
  EFI_STATUS Status = EFI_SUCCESS;
  UINT8      SlotIndex = 0;
  UINT8      SlaveAddress = 0;
  UINT8      Value = 0;
  BOOLEAN    Flag = FALSE;

  for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
    SlaveAddress = (DRAMSlaveAddress[SlotIndex] >> 1);
    Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, 0, &Value);
    if (EFI_ERROR (Status)) {
      continue;
    }

    if (Value == DDR5_SPD_FALG) {
      Flag = TRUE;
      break;
    }
  }

  return Flag;
}

EFI_STATUS
ReadSPD4Info ()
{
  EFI_STATUS  Status = EFI_SUCCESS;
  UINT8    SlotIndex = 0;
  UINT16   Index     = 0;
  UINT8    SlaveAddress = 0;
  UINT8    Value = 0;

  if (DDR4_SPD_SIZE <= 0 || DDR4_SPD_SIZE > 512) {
    return EFI_ABORTED;
  }

  for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
    SlaveAddress = (DRAMSlaveAddress[SlotIndex] >> 1);
    Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, 0, &Value);
    if (EFI_ERROR (Status)) {
      gDRAMInfo[SlotIndex].DRAMPresent = 0x00;
      continue;
    }

    gDRAMInfo[SlotIndex].DRAMPresent = 0x01;

    if (DDR4_SPD_SIZE <= 0x100) {
      WriteByteBySmbus ((0x6C >> 1), SMBUS_HOST_CTL_BYTE_DATA, 0, 0);
      for (Index = 0; Index < DDR4_SPD_SIZE; Index++) {
        Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, (UINT8)Index, &Value);
        if (!EFI_ERROR (Status)) {
          gDRAMInfo[SlotIndex].SPDData[Index] = Value;
        }
      }
    } else {
      WriteByteBySmbus ((0x6C >> 1), SMBUS_HOST_CTL_BYTE_DATA, 0, 0);
      for (Index = 0; Index < 0x100; Index++) {
        Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, (UINT8)Index, &Value);
        if (!EFI_ERROR (Status)) {
          gDRAMInfo[SlotIndex].SPDData[Index] = Value;
        }
      }

      WriteByteBySmbus ((0x6E >> 1), SMBUS_HOST_CTL_BYTE_DATA, 0, 0);
      Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, 0, &Value);
      if (EFI_ERROR (Status)) {
        continue;
      }
      
      for (Index = 0; Index < 0x100; Index++) {
        Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, (UINT8)Index, &Value);
        if (!EFI_ERROR (Status)) {
          gDRAMInfo[SlotIndex].SPDData[0x100 + Index] = Value;
        }
      }
    }
  }

#if 1
  {
    UINT16  RowIndex    = 0;
    UINT16  ColumnIndex = 0;
    UINT16  RowCount    = 0;
    UINT16  ColumnCount = 0;

    for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
      if (gDRAMInfo[SlotIndex].DRAMPresent == 0x00) {
        continue;
      }

      if (0 < DDR4_SPD_SIZE) {
        Print (L"------------\n");
        Print (L"Slot: 0x%02x |\n", DRAMSlaveAddress[SlotIndex]);
        Print (L"           ---------------------------------------------------------------------------\n");
        Print (L"        00   01   02   03   04   05   06   07   08   09   0A   0B   0C   0D   0E   0F\n");
        if ((DDR4_SPD_SIZE / 16) == 0) {
          Print (L"0x%04x ", (DDR4_SPD_SIZE / 16));
          for (ColumnIndex = 0; ColumnIndex < DDR4_SPD_SIZE; ColumnIndex++) {
            Print (L"0x%02x ", gDRAMInfo[SlotIndex].SPDData[ColumnIndex]);
          }
          Print (L"\n");
        } else {
          RowCount = (DDR4_SPD_SIZE / 16);
          for (RowIndex = 0; RowIndex <= RowCount; RowIndex++) {
            ColumnCount = (RowIndex == RowCount) ? (DDR4_SPD_SIZE - (RowCount * 16)) : 16;
            if (ColumnCount) {
              Print (L"0x%04x ", RowIndex * 16);
            }
            for (ColumnIndex = 0; ColumnIndex < ColumnCount; ColumnIndex++) {
              Print (L"0x%02x ", gDRAMInfo[SlotIndex].SPDData[RowIndex * 16 + ColumnIndex]);
            }
            if (ColumnCount) {
              Print (L"\n");
            }
          }
        }
        Print (L"--------------------------------------------------------------------------------------\n");
      }
      Print (L"\n");
    }
  }
#endif

  return Status;
}

VOID
PrintSPD4Info ()
{
  UINT8    SlotIndex = 0;

  for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
    if (gDRAMInfo[SlotIndex].DRAMPresent) {
      GetDRAMCapacity (gDRAMInfo[SlotIndex].SPDData[0x04], gDRAMInfo[SlotIndex].SPDData[0x06], gDRAMInfo[SlotIndex].SPDData[0x0C], gDRAMInfo[SlotIndex].SPDData[0x0D]);
      GetDRAMSpeed (gDRAMInfo[SlotIndex].SPDData[0x11], gDRAMInfo[SlotIndex].SPDData[0x12], gDRAMInfo[SlotIndex].SPDData[0x7D]);
      Print (L"Memory Vendor ManufacturerString: %a\n", Jep106GetManufacturerString ((gDRAMInfo[SlotIndex].SPDData[0x140] & 0x7F), (gDRAMInfo[SlotIndex].SPDData[0x141] & 0x7F)));
      Print (L"Module Serial Number: 0x%02x%02x%02x%02x\n", 
                gDRAMInfo[SlotIndex].SPDData[0x145], 
                gDRAMInfo[SlotIndex].SPDData[0x146], 
                gDRAMInfo[SlotIndex].SPDData[0x147], 
                gDRAMInfo[SlotIndex].SPDData[0x148]
                );
      Print (L"Module Part Number: %a\n", (UINT8 *)&(gDRAMInfo[SlotIndex].SPDData[0x149]));
      Print (L"SDRAM Vendor ManufacturerString: %a\n", GetDRAMKELIString (gDRAMInfo[SlotIndex].SPDData[0x15F], gDRAMInfo[SlotIndex].SPDData[0x15E]));
      Print (L"CheckSum: 0x%04x\n", CalculateDdr4DimmChecksum (gDRAMInfo[SlotIndex].SPDData));
      Print (L"CL: %d \n", GetCL (gDRAMInfo[SlotIndex].SPDData, 0x11, 0x12, 0x7D, 0x18, 0x7B));
    }
  }
}


EFI_STATUS
ReadSPD5Info ()
{
  EFI_STATUS  Status = EFI_SUCCESS;
  UINT8    SlotIndex = 0;
  UINT16   Index     = 0;
  UINT8    SlaveAddress = 0;
  UINT8    Value = 0;
  UINT8    Page  = 0;

  if (DDR5_SPD_SIZE <= 0 || DDR5_SPD_SIZE > 1024) {
    return EFI_ABORTED;
  }

  for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++ ) {
    SlaveAddress = DRAMSlaveAddress[SlotIndex] >> 1;
    Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, 0, &Value);
    if (EFI_ERROR (Status)) {
      gDDR5DRAMInfo[SlotIndex].DRAMPresent = 0x00;
      continue;
    }

    gDDR5DRAMInfo[SlotIndex].DRAMPresent = 0x01;
    for (Page = 0; Page < (DDR5_SPD_SIZE / 0x80); Page++) {
      WriteByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, DDR5_MR11, Page);
      for (Index = 0; Index < 0x80; Index++) {
        Status = ReadByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, (UINT8)(Index | BIT7), &Value);
        if (!EFI_ERROR (Status)) {
          gDDR5DRAMInfo[SlotIndex].SPDData[(Page * 0x80) + Index] = Value;
        }
      }
    }
    WriteByteBySmbus (SlaveAddress, SMBUS_HOST_CTL_BYTE_DATA, DDR5_MR11, 0);
  }


#if 1
  {
    UINT16  RowIndex    = 0;
    UINT16  ColumnIndex = 0;
    UINT16  RowCount    = 0;
    UINT16  ColumnCount = 0;

    for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
      if (gDDR5DRAMInfo[SlotIndex].DRAMPresent == 0x00) {
        continue;
      }

      if (0 < DDR5_SPD_SIZE) {
        Print (L"------------\n");
        Print (L"Slot: 0x%02x |\n", DRAMSlaveAddress[SlotIndex]);
        Print (L"           ---------------------------------------------------------------------------\n");
        Print (L"        00   01   02   03   04   05   06   07   08   09   0A   0B   0C   0D   0E   0F\n");
        if ((DDR5_SPD_SIZE / 16) == 0) {
          Print (L"0x%04x ", (DDR5_SPD_SIZE / 16));
          for (ColumnIndex = 0; ColumnIndex < DDR5_SPD_SIZE; ColumnIndex++) {
            Print (L"0x%02x ", gDDR5DRAMInfo[SlotIndex].SPDData[ColumnIndex]);
          }
          Print (L"\n");
        } else {
          RowCount = (DDR5_SPD_SIZE / 16);
          for (RowIndex = 0; RowIndex <= RowCount; RowIndex++) {
            ColumnCount = (RowIndex == RowCount) ? (DDR5_SPD_SIZE - (RowCount * 16)) : 16;
            if (ColumnCount) {
              Print (L"0x%04x ", RowIndex * 16);
            }
            for (ColumnIndex = 0; ColumnIndex < ColumnCount; ColumnIndex++) {
              Print (L"0x%02x ", gDDR5DRAMInfo[SlotIndex].SPDData[RowIndex * 16 + ColumnIndex]);
            }
            if (ColumnCount) {
              Print (L"\n");
            }
          }
        }
        Print (L"--------------------------------------------------------------------------------------\n");
      }
      Print (L"\n");
    }
  }
#endif

  return Status;
}

VOID
PrintSPD5Info ()
{
  UINT8    SlotIndex = 0;

  for (SlotIndex = 0; SlotIndex < MAX_DRAM_SLOT_NUM; SlotIndex++) {
    if (gDDR5DRAMInfo[SlotIndex].DRAMPresent) {
      Print (L"Memory Vendor ManufacturerString: %a\n", Jep106GetManufacturerString ((gDDR5DRAMInfo[SlotIndex].SPDData[0x200] & 0x7F), (gDDR5DRAMInfo[SlotIndex].SPDData[0x201] & 0x7F)));
    }
  }
}

/*
0x04 SDRAM Density and Banks 
0x06 SDRAM Package Type
0x0C Module Organization
0x0D Module Memory Bus Width
*/
VOID
GetDRAMCapacity (
  IN UINT8  DensityAndBanks,
  IN UINT8  PackageType,
  IN UINT8  ModuleOrganization,
  IN UINT8  ModuleBusWidth
  )
{
  UINTN  SDRAMCapacity   = 0;
  UINT8  PrimaryBusWidth = 0;
  UINT8  SDRAMWidth      = 0;
  UINT8  PackageRanksPDIMM = 0;
  UINT8  LogicalRanksPerDimm = 0;

  switch ((DensityAndBanks & 0x0F)) {
    case 0:
      SDRAMCapacity = SIZE_1MB * 256; // 256Mb
      break;
    case 1:
      SDRAMCapacity = SIZE_1MB * 512; // 512Mb
      break;
    case 2:
      SDRAMCapacity = SIZE_1GB; // 1Gb
      break;
    case 3:
      SDRAMCapacity = (UINTN)((UINTN)SIZE_1GB * (UINTN)2); // 2Gb
      break;
    case 4:
      SDRAMCapacity = (UINTN)((UINTN)SIZE_1GB * (UINTN)4); // 4Gb
      break;
    case 5:
      SDRAMCapacity = (UINTN)((UINTN)SIZE_1GB * (UINTN)8); // 8Gb
      break;
    case 6:
      SDRAMCapacity = (UINTN)((UINTN)SIZE_1GB * (UINTN)16); // 16Gb
      break;
    case 7:
      SDRAMCapacity = (UINTN)((UINTN)SIZE_1GB * (UINTN)32); // 32Gb
      break;
    default:
      Print (L"SDRAMCapacity Value Error\n");
      return;
  }

  switch (ModuleBusWidth & 0x07) {
    case 0:
      PrimaryBusWidth = 8; // 8 bits
      break;
    case 1:
      PrimaryBusWidth = 16; // 16 bits
      break;
    case 2:
      PrimaryBusWidth = 32; // 32 bits
      break;
    case 3:
      PrimaryBusWidth = 64; // 64 bits
      break;
    default:
      Print (L"PrimaryBusWidth Value Error\n");
      return;
  }

  switch (ModuleOrganization & 0x07) {
    case 0:
      SDRAMWidth = 4; // 4 bits
      break;
    case 1:
      SDRAMWidth = 8; // 8 bits
      break;
    case 2:
      SDRAMWidth = 16; // 16 bits
      break;
    case 3:
      SDRAMWidth = 32; // 32 bits
      break;
    default:
      Print (L"SDRAMWidth Value Error\n");
      return;
  }

  switch ((ModuleOrganization >> 3) & 0x07) {
    case 0:
      PackageRanksPDIMM = 1; // 1 Package Rank
      break;
    case 1:
      PackageRanksPDIMM = 2; // 2 Package Ranks
      break;
    case 2:
      PackageRanksPDIMM = 3; // 3 Package Ranks
      break;
    case 3:
      PackageRanksPDIMM = 4; // 4 Package Ranks
      break;
    default:
      Print (L"PackageRanksPDIMM Value Error\n");
      return;
  }

  if ((PackageType & 0x03) == 0x02) { // 3DS
    LogicalRanksPerDimm = PackageRanksPDIMM * (((PackageType >> 4) & 0x07) + 1);
  } else {
    LogicalRanksPerDimm = PackageRanksPDIMM * 1;
  }

  /*
  Print (L"SDRAMCapacity: 0x%llx \nPrimaryBusWidth: 0x%x \nSDRAMWidth: 0x%x \nPackageRanksPDIMM: 0x%x \nLogicalRanksPerDimm: 0x%x \n", 
            SDRAMCapacity,
            PrimaryBusWidth,
            SDRAMWidth,
            PackageRanksPDIMM,
            LogicalRanksPerDimm
            );
  */

  SDRAMCapacity = (((SDRAMCapacity / 8) * PrimaryBusWidth) / SDRAMWidth) * LogicalRanksPerDimm;

  Print (L"DIMM Size %d GB\n", (SDRAMCapacity / (UINTN)SIZE_1GB));
}

UINT16
GetDRAMSpeed (
  IN UINT8 TimeBases,
  IN INT8 MinCycleTime,
  IN INT8 FindMinCycleTime
  )
{
  INTN   Tck =0;
  UINT16 Speed = 0;

  Tck = MinCycleTime * 125 + FindMinCycleTime * 1;
  
  Speed = (UINT16)(1000000 / Tck * 2);

  Print (L"Speed %d MHz\n", Speed);

  return Speed;
}

UINT16
GetCL(
  IN UINT8 *DataBuffer,
  IN UINT8 TimeBasesIndex,
  IN UINT8 MinCycleTimeIndex,
  IN UINT8 FindMinCycleTimeIndex,
  IN UINT8 MCASLatencyTimeIndex,
  IN UINT8 FindMinCASLatencyTimeIndex
  )
{
  UINT8 TimeBases = 0;
  INT8 MinCycleTime = 0;
  INT8 FindMinCycleTime = 0;
  INT8 MCASLatencyTime = 0;
  INT8 FindMinCASLatencyTime = 0;
  INTN   Tcl = 0;
  INTN   Tck = 0;

  TimeBases = DataBuffer[TimeBasesIndex];
  MinCycleTime = DataBuffer[MinCycleTimeIndex];
  FindMinCycleTime = DataBuffer[FindMinCycleTimeIndex];
  MCASLatencyTime = DataBuffer[MCASLatencyTimeIndex];
  FindMinCASLatencyTime = DataBuffer[FindMinCASLatencyTimeIndex];

  if ((TimeBases & 0x0F) == 0) {
    Tck = (MinCycleTime * 125 * 10) + (FindMinCycleTime * 10);
    Tcl = (MCASLatencyTime * 125 * 10) + (FindMinCASLatencyTime * 10);
    Tcl = (Tcl % Tck) ? (Tcl / Tck) + 1 : (Tcl / Tck);
  }

  return (UINT16)Tcl;
}


UINT16 
Crc16 (
  UINT8 *ptr,
  INT16 count
  )
{
  UINT16 crc, i;
  crc = 0;

  while (--count >= 0) {
    crc = crc ^ (int)*ptr++ << 8;
    for (i = 0; i < 8; ++i){
      if (crc & 0x8000) {
        crc = crc << 1 ^ 0x1021;
      } else {
        crc = crc << 1;
      }
    }
  }
  
  return (crc & 0xFFFF);
}


UINT16
CalculateDdr4DimmChecksum (
  IN UINT8 *DataBuffer
  )
{
  UINT16 checksum = 0;
  UINT16 wIndex = 0,i = 0;
  UINT8  bChecksum[512];
  // calculate checksum with customized algorithm
  // TODO, we should use CRC16 to get the checksum
  checksum = 0;

  // CRC for Base Configuration Section
  bChecksum[wIndex++] = DataBuffer[126];
  bChecksum[wIndex++] = DataBuffer[127];

  // Module Manufacturer ID Code
  bChecksum[wIndex++] = DataBuffer[320];
  bChecksum[wIndex++] = DataBuffer[321];

  // Module Part Number
  for(i = 329;i <= 348;i++) {
    bChecksum[wIndex++] = DataBuffer[i];
  }

  // DRAM Manufacturer ID Code
  bChecksum[wIndex++] = DataBuffer[350];
  bChecksum[wIndex++] = DataBuffer[351];

  checksum = Crc16(bChecksum,wIndex);

  return checksum;
}